import os
from typing import List

from langchain.chains import LLMChain
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from backend.app.core.conf import settings
from backend.app.schemas.device import DeviceAttributeListSchema

os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY
os.environ["OPENAI_API_BASE"] = settings.OPENAI_API_BASE


class DataCleanChain:
    _llm = None
    _parser = None
    _output_key = "result"

    def __init__(self):
        self._llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
        self._parser = PydanticOutputParser(pydantic_object=DeviceAttributeListSchema)

    @property
    def template(self):
        return """您是一个音视频设备专家。您将收到一个设备的属性。您的目标是分类归纳。您需要将属性中符号描述转为文字描述。
{format_instructions}
设备属性将以(```)包裹.
```{query}```
最终答案："""

    def get_prompt(self):
        return PromptTemplate(
            template=self.template,
            input_variables=["query"],
            partial_variables={"format_instructions": self._parser.get_format_instructions()}
        )

    def get_chain(self):
        prompt = self.get_prompt()
        return LLMChain(
            llm=self._llm,
            prompt=prompt,
            output_key=self._output_key
        )

    def parser_output(self, output):
        json_str = output[self._output_key]
        parser = self._parser
        try:
            res = parser.parse(json_str)
        except Exception as e:
            print(e)
            try:
                new_parser = OutputFixingParser.from_llm(parser=parser, llm=self._llm)
                res = new_parser.parse(json_str)
            except Exception as e:
                print(e)
                res = None
        return res

    @staticmethod
    def format_attribute(meters) -> str:
        """还原设备属性"""
        lists = []
        for m in sorted(meters, key=lambda x: x.attribute):
            if m.value:
                v = f"{m.attribute}:{m.value}"
            else:
                v = m.attribute
            lists.append(v)
        return ";".join(lists)

    def clean_data(self, attribute: str):
        """清洗设备属性"""
        chain = self.get_chain()
        output = chain.invoke(input=dict(query=attribute))
        obj = self.parser_output(output)
        res = self.format_attribute(obj.lists)
        return res

    async def aclean_data(self, attribute: str) -> str:
        chain = self.get_chain()
        output = await chain.ainvoke(input=dict(query=attribute))
        obj = self.parser_output(output)
        res = self.format_attribute(obj.lists)
        return res

    async def abatch_clean_data(self, meter_list: List[str]) -> List[str]:
        chain = self.get_chain()
        lists = []
        input_list = [dict(query=m) for m in meter_list]
        output_list = await chain.abatch(inputs=input_list)
        for index, output in enumerate(output_list):
            obj = self.parser_output(output)
            v = meter_list[index]
            if obj and obj.lists:
                v = self.format_attribute(obj.lists)
            lists.append(v)
        return lists

